#!/bin/bash

# ============================================================================ #
# Copyright (c) 2022 - 2025 NVIDIA Corporation & Affiliates.                   #
# All rights reserved.                                                         #
#                                                                              #
# This source code and the accompanying materials are made available under     #
# the terms of the Apache License 2.0 which accompanies this distribution.     #
# ============================================================================ #

# This script allows us to allocate a single GPU 
# to one of the MPI ranks.
numGPUS=$(nvidia-smi --list-gpus | wc -l)
echo $numGPUS
CUDA_VISIBLE_DEVICES=$(($OMPI_COMM_WORLD_RANK % $numGPUS)) @Python_EXECUTABLE@ $1
